"""
Basic image transformation functions
Images are always in PIL.Image type
"""
import torch
import numpy as np
from PIL import Image
import torchvision.transforms.functional as FT


class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, boxes=None, labels=None):
        for t in self.transforms:
            image, boxes, labels = t(image, boxes, labels)
        return image, boxes, labels


# class ToTensor(object):
#     """
#     geometric and photometric transformations are applied before ToTensor
#     """
#     def __call__(self, image, boxes=None, labels=None):
#         """
#         image: PIL Image with mode=RGB
#         boxes: numpy array of shape [nobj, 4], in percent coords
#         labels: numpy array of shape [nobj,]
#         """
#         image = torch.from_numpy(np.array(image))
#         image = image.permute((2, 0, 1)).contiguous()
#         image = image.float().div(255)
#         if boxes is not None:
#             boxes = torch.FloatTensor(boxes)
#         if labels is not None:
#             labels = torch.LongTensor(labels)

#         return image, boxes, labels


# class Normalize(object):
#     def __init__(self, mean=0.0, std=1.0):
#         self.mean = mean
#         self.std = std
    
#     def __call__(self, image, boxes=None, labels=None, inplace=False):
#         if not inplace:
#             image = image.clone()
        
#         dtype = image.dtype
#         mean = torch.as_tensor(self.mean, dtype=dtype, device=image.device)
#         std = torch.as_tensor(self.std, dtype=dtype, device=image.device)
#         # if mean.ndim == 1:
#         mean = mean[:, None, None]
#         # if std.ndim == 1:
#         std = std[:, None, None]
#         image.sub_(mean).div_(std)
#         return image, boxes, labels


class NullTransform(object):
    def __call__(self, image, boxes, labels):
        return image, boxes, labels


class ToTensor(object):
    def __call__(self, image, boxes, labels):
        image = FT.to_tensor(image)
        return image, boxes, labels


class Normalize(object):
    """
    normalize is applied after totensor when image is an torch.Tensor
    """
    def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        self.mean = mean
        self.std = std
    
    def __call__(self, image, boxes, labels):
        image = FT.normalize(image, mean=self.mean, std=self.std)
        return image, boxes, labels


class ToOpenCV(object):
    def __call__(self, image, boxes=None, labels=None):
        """
        image: PIL Image
        """
        image = np.array(image)
        return image, boxes, labels


class ToPIL(object):
    def __call__(self, image, boxes=None, labels=None):
        image = PIL.Image.fromarray(image)
        return image, boxes, labels
